# 导包

import os,time
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
from tensorflow.keras.models import save_model

# from nets.conv_net import ConvModel
# from utils.data_generator import train_val_generator
# from utils.image_plot import plot_images

from nets.conv_net import ConvModel
from utils.data_generator import train_val_generator
from utils.image_plot import plot_images

# import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '2'

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = "true"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = '0'


# 读取训练集
train_gen = train_val_generator(
    data_dir='../dataset/train',
    target_size=(64, 64),
    batch_size=32,
    class_mode='categorical',
    subset="training"
)

# 读取验证集
val_gen = train_val_generator(
    data_dir='../dataset/train',
    target_size=(64, 64),
    batch_size=32,
    class_mode='categorical',
    subset="validation"
)

# 获取标签名列表
class_names = list(train_gen.class_indices.keys())

# ImageDataGenerator的返回结果是迭代器，
# 取15张训练集图片进行查看
train_batch, train_label_batch = train_gen.next()
plot_images(train_batch, train_label_batch, class_names)

# 取15张验证集图片查看
val_batch, val_label_batch = val_gen.next()
plot_images(val_batch, val_label_batch, class_names)

"""
模型设置 tf.keras.Sequential.compile
参数：
-loss
-optimizer "sgd"
-metrics   "accuracy"

"""

model = ConvModel()


# 设置损失


model.compile(
    loss='categorical_crossentropy',
    optimizer=tf.keras.optimizers.SGD(
        learning_rate=0.001
    ),
    metrics=['accuracy']
)

"""
 模型训练 tf.keras.Sequential.fit
 参数：
 -x
 -steps_per_epoch: 图片总量/batch_size得到
 -epochs:
 -validation_data:
 -validation_steps:验证集跑多少步来计算模型的评价指标， 一步会读取batch_size张图片，所以一共验证validation_steps*batch_size张图片
 -shuffle: 默认为True
 
 返回： History.history 记录每一轮训练集和验证集的损失函数和评价指标
 
 # 237 = 9469/32*0.8
 # 60 = 9469/32*0.2
"""

# 模型训练
history = model.fit(
    x=train_gen,
    steps_per_epoch=237,
    epochs=100,
    validation_data=val_gen,
    validation_steps=60,
    shuffle=True
)



# 画图查看history的变化趋势

pd.DataFrame(history.history).plot(figsize=(8, 5))
plt.grid(True)
plt.xlabel('epochs')
plt.show()



"""
模型保存 tf.keras.models.save_model
参数：
-model:
-filepath
"""

# 创建保存路径
model_name = "model-" + time.strftime("%Y-%m-%d-%H-%M-%S")
model_path = os.path.join('..', "models", model_name)
if not os.path.exists(model_path):
    os.makedirs(model_path)

# 模型保存
save_model(model=model, filepath=model_path)

model.summary()
